#!/usr/bin/env python3
"""
测试肤色融合功能
"""

import sys
import os
import numpy as np
from PIL import Image
import cv2

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from image_utils.skin_color_fusion import skin_color_fusion_pipeline, extract_average_skin_color, apply_grayscale_based_skin_color

def create_test_images():
    """创建测试图像"""
    # 创建参考图像（带肤色的角色）
    reference_img = np.zeros((512, 512, 3), dtype=np.uint8)
    reference_img[100:400, 150:350] = [220, 180, 160]  # 肤色区域
    reference_img[200:300, 200:300] = [255, 255, 255]  # 衣服区域
    
    # 创建参考图肤色mask
    reference_mask = np.zeros((512, 512), dtype=np.uint8)
    reference_mask[100:400, 150:350] = 255  # 肤色区域
    
    # 创建带灰度的白模图像（模拟真实白模）
    white_model = np.ones((512, 512, 3), dtype=np.uint8) * 240  # 白色背景
    
    # 创建渐变的身体区域，模拟光照效果
    for y in range(100, 400):
        for x in range(150, 350):
            # 中心最亮，边缘较暗
            center_x, center_y = 250, 250
            distance = np.sqrt((x - center_x)**2 + (y - center_y)**2)
            max_distance = np.sqrt(100**2 + 100**2)
            intensity = 1.0 - (distance / max_distance) * 0.4  # 0.6到1.0的强度
            intensity = max(0.3, intensity)  # 最低0.3的强度
            gray_value = int(255 * intensity)
            white_model[y, x] = [gray_value, gray_value, gray_value]
    
    # 创建白模mask
    white_mask = np.zeros((512, 512), dtype=np.uint8)
    white_mask[100:400, 150:350] = 255  # 身体区域
    
    return reference_img, reference_mask, white_model, white_mask

def test_skin_fusion():
    """测试肤色融合功能"""
    print("创建测试图像...")
    ref_img, ref_mask, white_img, white_mask = create_test_images()
    
    print("执行肤色融合...")
    fused_img, fusion_mask = skin_color_fusion_pipeline(
        ref_img, ref_mask, white_img, white_mask,
        texture_strength=0.2
    )
    
    # 测试平均肤色提取
    print("测试平均肤色提取...")
    avg_skin_color = extract_average_skin_color(ref_img, ref_mask)
    print(f"提取的平均肤色: {avg_skin_color}")
    
    # 测试基于灰度的融合
    print("测试基于灰度的融合...")
    grayscale_fused = apply_grayscale_based_skin_color(white_img, white_mask, avg_skin_color)
    
    print("保存结果图像...")
    # 保存结果
    Image.fromarray(ref_img).save("test_reference.png")
    Image.fromarray(ref_mask).save("test_reference_mask.png")
    Image.fromarray(white_img).save("test_white_model.png")
    Image.fromarray(white_mask).save("test_white_mask.png")
    Image.fromarray(fused_img).save("test_fused_result.png")
    Image.fromarray(fusion_mask).save("test_fusion_mask.png")
    Image.fromarray(grayscale_fused).save("test_grayscale_fused.png")
    
    print("测试完成！结果已保存到当前目录")
    print("- test_reference.png: 参考图像")
    print("- test_reference_mask.png: 参考图肤色mask")
    print("- test_white_model.png: 白模图像（带灰度渐变）")
    print("- test_white_mask.png: 白模mask")
    print("- test_fused_result.png: 完整融合结果")
    print("- test_grayscale_fused.png: 基于灰度的融合结果")
    print("- test_fusion_mask.png: 融合区域mask")

if __name__ == "__main__":
    test_skin_fusion()
